library(tidyverse)

source("Code/helper_functions.R")

# Function to post-process the variational Bayes fit and extract municipality-level estimates
post_process_vb_fit <- function(municipal_level_vb_fit, input_data) {
  # Aggregate input data to municipality level and prepare data for Stan
  stan_data_municipality_level <- input_data %>%
    group_by(covid, Municipality_code, Municipality_idx, Municipality_label) %>%
    summarise(Population = sum(Population), NbDeath = sum(NbDeath)) %>%
    create_stan_data()
  
  # Extract posterior estimates from the fit and convert to interpretable quantities
  municipality_level_vb_estimates <- tibble(
    excess_mortality_parameter = municipal_level_vb_fit@sim$est$log_lambda_c, # Extract excess mortality parameter
    baseline_mortality_parameter = municipal_level_vb_fit@sim$est$log_lambda0 # Extract baseline mortality parameter
  ) %>%
    rowid_to_column(var = "municipality_idx") %>% # Add municipality index as a column
    mutate(Municipality_code = stan_data_municipality_level$Municipality_code_converter(municipality_idx)) %>% # Convert index to code
    mutate(excess_mortality_factor = exp(excess_mortality_parameter)) %>% # Transform to excess mortality factor
    mutate(baseline_mortality = exp(baseline_mortality_parameter)) %>% # Transform to baseline mortality
    select(Municipality_code, excess_mortality_factor, baseline_mortality) # Select relevant columns
  
  return(municipality_level_vb_estimates) # Return the estimates
}

#######################################
municipal_level_vb_fit <- read_rds("Data/Saves/electoral_stability_municipality_level_vb_no_2020_2.rds")

full_data <- read_rds("Data/Saves/full_mortality_data_2015_2021_no_2019_second_semester_2020_first_semester.rds")

municipality_level_vb_estimates <- post_process_vb_fit(municipal_level_vb_fit, full_data)


if (!dir.exists("Results")) dir.create("Results")

municipality_level_vb_estimates %>%
  rename(
    excess_mortality_factor_vb_remaining_waves = excess_mortality_factor,
    baseline_mortality_vb_remaining_waves = baseline_mortality
  ) %>%
  write.csv(x = ., file = "Results/municipality_excess_mortality_vb_electoral_stability_2024_11_27.csv", row.names = F)
